import h5py
import numpy as np

from VGG13_v15 import config as vgg13_config
from VGG13_v15 import image as IM

rng = np.random.RandomState(vgg13_config.seed)


class Generator():
    def __init__(self, hdf5_path, candidates_path, valid_fold, include_test):
        candidates_tuple = np.load(candidates_path)
        self.train_candidates = list()
        self.val_candidates = list()

        self.candidates_dataset = h5py.File(hdf5_path, 'r')
        self.candidates_pool = dict()
        self.centers = np.zeros((3, 2048))
        self.lr = vgg13_config.center_lr

        cnt_malign = 0
        cnt_benign = 0
        cnt_neg = 0

        for candidate in candidates_tuple:
            if candidate['fold'] == valid_fold:
                if candidate['source'] == 'kaggle_beni':
                    if candidate['prob'] >= 0.65:
                        self.val_candidates.append((candidate, 1))
                elif candidate['source'] == 'spie':
                    if candidate['malign'] == 1:
                        self.val_candidates.append((candidate, 0))
                    elif candidate['malign'] == 0:
                        self.val_candidates.append((candidate, 1))
                elif candidate['source'] == 'lidc':
                    if candidate['polarity'] == 0:
                        self.val_candidates.append((candidate, 1))
                    elif candidate['polarity'] == 1:
                        if candidate['malign'] >= 4:
                            self.val_candidates.append((candidate, 0))
                        elif candidate['malign'] < 3:
                            self.val_candidates.append((candidate, 1))
                    elif candidate['prob'] >= 0.65:
                        self.val_candidates.append((candidate, 1))
                elif candidate['source'] == 'kaggle_mali' or (candidate['source'] == 'kaggle_testset' and include_test):
                    if candidate['polarity'] == 1:
                        if candidate['malign'] >= 7:
                            self.val_candidates.append((candidate, 0))
                        elif candidate['malign'] < 5:
                            self.val_candidates.append((candidate, 1))
                    elif candidate['prob'] >= 0.65:
                        self.val_candidates.append((candidate, 1))
            else:
                if candidate['source'] == 'kaggle_beni':
                    if candidate['prob'] >= 0.9:
                        self.train_candidates.append((candidate, 1))
                    elif 0.65 <= candidate['prob'] <= 0.8:
                        self.train_candidates.append((candidate, 2))
                elif candidate['source'] == 'spie':
                    if candidate['malign'] == 1:
                        self.train_candidates.append((candidate, 0))
                    elif candidate['malign'] == 0:
                        self.train_candidates.append((candidate, 1))
                elif candidate['source'] == 'lidc':
                    if candidate['polarity'] == 0:
                        self.train_candidates.append((candidate, 1))
                    elif candidate['polarity'] == 1:
                        if candidate['malign'] >= 4:
                            self.train_candidates.append((candidate, 0))
                        elif candidate['malign'] < 3:
                            self.train_candidates.append((candidate, 1))
                    elif candidate['prob'] >= 0.65:
                        self.train_candidates.append((candidate, 2))
                elif candidate['source'] == 'kaggle_mali' or (candidate['source'] == 'kaggle_testset' and include_test):
                    if candidate['polarity'] == 1:
                        if candidate['malign'] >= 7:
                            self.train_candidates.append((candidate, 0))
                        elif candidate['malign'] < 5:
                            self.train_candidates.append((candidate, 1))
                    elif candidate['prob'] >= 0.9:
                        self.train_candidates.append((candidate, 1))

        for _, label_type in self.train_candidates:
            if label_type == 0:
                cnt_malign += 1
            elif label_type == 1:
                cnt_benign += 1
            else:
                cnt_neg += 1

        print('Train: malign=%d, benign=%d, neg=%d, overall=%d' %(cnt_malign, cnt_benign, cnt_neg, len(self.train_candidates)))

    def update_center(self, features, labels):
        transpose_labels = np.transpose(labels)
        delta_centers = (np.matmul(labels, self.centers) - features)
        delta_centers = np.matmul(transpose_labels, delta_centers)
        self.centers -= self.lr * delta_centers/ (1.0 + np.sum(transpose_labels, axis=1, keepdims=True))

    def get_subvol(self, candidate):
        path = candidate['path']
        subvol = self.candidates_pool.get(path)
        if subvol is None:
            subvol = self.candidates_dataset[path][()]
            self.candidates_pool[path] = subvol

        return subvol

    def train_generate(self, batch_size=1):
        subvols = list()
        labels = list()
        centers = list()
        while True:
            rng.shuffle(self.train_candidates)

            for candidate, label_type in self.train_candidates:
                subvol = self.get_subvol(candidate)
                subvol = random_crop(IM.normalize(subvol,-1024,800))
                subvol = random_augmentation(subvol, 0.5)
                subvols.append(subvol)
                label = np.zeros(3)
                label[label_type] = 1
                labels.append(label)
                centers.append(self.centers[label_type, :])

                if len(subvols) == batch_size:
                    yield np.expand_dims(np.array(subvols), -1), np.array(labels), np.array(centers)
                    subvols.clear()
                    labels.clear()
                    centers.clear()

    def generate_val_labels(self):
        self.val_pos = 0
        self.val_neg = 0
        val_labels = list()
        for candidate, label_type in self.val_candidates:
            if label_type == 0:
                self.val_pos += 1
                val_labels.append([1, 0])
            elif label_type == 1:
                self.val_neg += 1
                val_labels.append([0, 1])

        self.val_labels = np.asarray(val_labels)
        print('validate: pos=%d, neg=%d, overall=%d' %(self.val_pos, self.val_neg, len(self.val_candidates)))

    def val_generate(self, batch_size=1):
        subvols = list()
        while True:
            for candidate, label_type in self.val_candidates:
                subvol = self.get_subvol(candidate)
                subvol = middle_crop(IM.normalize(subvol, -1024, 800))
                subvols.append(subvol)

                if len(subvols) == batch_size:
                    yield np.expand_dims(np.array(subvols), -1)
                    subvols.clear()

            if len(subvols) != 0:
                yield np.expand_dims(np.array(subvols), -1)
                subvols.clear()


class Valid_Generator():
    def __init__(self, hdf5_path, candidates_path, valid_fold, include_test):
        candidates_tuple = np.load(candidates_path)
        self.val_candidates = list()
        self.candidates_dataset = h5py.File(hdf5_path, 'r')

        for candidate in candidates_tuple:
            if candidate['fold'] == valid_fold \
                    and (candidate['source'] != 'kaggle_testset' or include_test) \
                    and candidate['prob'] >= 0.65:
                self.val_candidates.append(candidate)

    def val_generate(self, batch_size=1, reverse=False):
        subvols = list()
        while True:
            for candidate in self.val_candidates:
                subvol = self.candidates_dataset[candidate['path']][()]

                if reverse:
                    subvol = subvol[:,:,::-1]

                subvol = middle_crop(IM.normalize(subvol, -1024, 800))
                subvols.append(subvol)

                if len(subvols) == batch_size:
                    yield np.expand_dims(np.array(subvols), -1)
                    subvols.clear()

            if len(subvols) != 0:
                yield np.expand_dims(np.array(subvols), -1)
                subvols.clear()


class Test_Generator(object):
    def __init__(self, hdf5_path, candidates_path):
        candidates_tuple = np.load(candidates_path)
        self.val_candidates = list()
        self.candidates_dataset = h5py.File(hdf5_path, 'r')

        for candidate in candidates_tuple:
            if candidate['prob'] >= 0.65:
                self.val_candidates.append(candidate)

    def val_generate(self, batch_size=1, reverse=False):
        subvols = list()
        while True:
            for candidate in self.val_candidates:
                subvol = self.candidates_dataset[candidate['path']][()]

                if reverse:
                    subvol = subvol[:,:,::-1]

                subvol = middle_crop(IM.normalize(subvol, -1024, 800))
                subvols.append(subvol)

                if len(subvols) == batch_size:
                    yield np.expand_dims(np.array(subvols), -1)
                    subvols.clear()

            if len(subvols) != 0:
                yield np.expand_dims(np.array(subvols), -1)
                subvols.clear()


def random_augmentation(x, prob):
    if rng.rand() < prob:
        for dim in range(x.ndim):
            if rng.rand() < 0.5:
                x = IM.flip_axis(x, dim)
        x = np.transpose(x, rng.shuffle(np.arange(x.ndim)))
    return x


def random_crop(x):
    crop_bbox = np.append([0, 0, 0], x.shape)
    crop_offset = rng.random_integers(0, 4, 3)
    crop_bbox[0:3] = crop_bbox[0:3] + crop_offset
    crop_bbox[3:6] = crop_bbox[3:6] - 4 + crop_offset
    x = IM.crop(x, crop_bbox)
    return x


def middle_crop(x):
    crop_bbox = np.append([0, 0, 0], x.shape)
    crop_bbox[0:3] = crop_bbox[0:3] + 2
    crop_bbox[3:6] = crop_bbox[3:6] - 2
    x = IM.crop(x, crop_bbox)
    return x


def tri2bi(tri_labels):
    bi_labels = np.zeros((tri_labels.shape[0],2))
    bi_labels[:, 0] = (tri_labels[:, 0] == 1)
    bi_labels[:, 1] = (tri_labels[:, 0] != 1)
    return bi_labels
